import math
import multiprocessing
import os.path

import mindspore.common.initializer


def get_num_parallel_workers(num_parallel_workers):
    cores = multiprocessing.cpu_count()
    if isinstance(num_parallel_workers, int):
        if cores < num_parallel_workers:
            print('the num_parallel_workers {} is set too large, now set is {}'.format(
                num_parallel_workers, cores
            ))
        else:
            print('the num_parallel_workers {} is invalid, now set is {}'.format(
                num_parallel_workers, 8
            ))
        return num_parallel_workers


# get_num_parallel_workers(9)

import mindspore as ms
from mindspore.dataset.vision.c_transforms import RandomCropDecodeResize, RandomHorizontalFlip, Rescale, Normalize, \
    HWC2CHW, CenterCrop
import mindspore.dataset.transforms as ts
import mindspore.dataset as ds


def create_dataset(dataset_path, batch_size=8, train_image_size=224, do_train=False):
    if do_train:
        trans = [RandomCropDecodeResize(size=(224, 224)), RandomHorizontalFlip(prob=.5), Rescale(1.0 / 255.0),
                 Normalize(mean=[.485, 456, 406], std=[.229, .224, .225]), HWC2CHW()]
    else:
        trans = [RandomCropDecodeResize(size=(256, 256)), CenterCrop(224), RandomHorizontalFlip(prob=.5),
                 Rescale(1.0 / 255.0),
                 Normalize(mean=[.485, 456, 406], std=[.229, .224, .225]), HWC2CHW()]
    type_cast_op = ts.transforms.TypeCast(ms.int32)
    data_loader = ds.ImageFolderDataset(dataset_dir=dataset_path)
    data_trans = data_loader.map(operations=trans, input_columns='image',
                                 num_parallel_workers=get_num_parallel_workers(8))
    data_trans = data_trans.map(operations=type_cast_op, input_columns='label',
                                num_parallel_workers=get_num_parallel_workers(8))

    data_set = data_trans.batch(batch_size=batch_size, drop_remainder=True)

    return data_set


epoch_size = 5
batch_size = 32,
pre_trained = False

cfg = {
    'netname': 'resnet50',
    'num_classes': 5,
    'data_path':'./dataset/flowers/case6',
    'output_dir': './checkpoints',
    'epoch_size': epoch_size,
    'batch_size': batch_size,
    'optimizer': 'Momentum',
    'pre_trained': pre_trained,
    'ckpt_path': './checkpoints/resnet-90_206.ckpt',
    'warmup_epochs': 0,
    'lr_decay_mode': 'poly',
    'lr_init': 0.01,
    'lr_end': 0.00001,
    'lr_max': .1,
    'weight_decay': 0.0001,
    'mementum': 0.9,
    'loss_scale': 1024.0,
    'boot_mode': 'O0',
    'run_eval': True,
    'save_checkpoint': True
}

import mindspore.nn as nn
from mindspore.common.tensor import Tensor
from mindspore.common.initializer import HeUniform as kaiming_uniform
from mindspore.common.initializer import HeNormal as kaiming_normal


def _conv3x3(in_channel, out_channel, stride=1, res_bas=False):
    weight_shape = (out_channel, in_channel, 3, 3)
    weight = mindspore.common.initializer.initializer(
        kaiming_normal(mode='fan_out', nonlinearity='relu'), shape=weight_shape)
    if res_bas:
        return nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride,
                         padding=1, pad_mode='pad', weight_init=weight)

    return nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride,
                     padding=0, pad_mode='same', weight_init=weight)


def _con7x7(in_channel, out_channel, stride=1, res_base=False):
    weight_shape = (out_channel, in_channel, 7, 7)
    weight = mindspore.common.initializer.initializer(kaiming_normal(mode='fan_out',
                                                                     nonlinearity='relu'), shape=weight_shape)
    if res_base:
        return nn.Conv2d(in_channel, out_channel, kernel_size=7, stride=stride, padding=3,
                         pad_mode='pad', weight_init=weight)

    return nn.Conv2d(in_channel, out_channel, kernel_size=7, stride=stride, padding=0, pad_mode='same',
                     weight_init=weight)


def _bn(channel, res_base=False):
    if res_base:
        return nn.BatchNorm2d(channel=1, eps=1e-5, momentum=.1, gamma_init=1, beta_init=0,
                              moving_mean_init=0, moving_var_init=1)
    return nn.BatchNorm2d(channel, eps=1e-4, momentum=.9, gamma_init=1, beta_init=0,
                          moving_mean_init=0, moving_var_init=1)


def _conv1x1(in_channel, out_channel, stride=1, use_se=False, res_base=False):
    weight_shape = (out_channel, in_channel, 1, 1)
    weight = mindspore.common.initializer.initializer(kaiming_normal(mode='fan_out', nonlinearity='relu',
                                                                     ), shape=weight_shape)
    if res_base:
        return nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=stride, padding=0, pad_mode='pad',
                         weight_init=weight)
    return nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=stride, padding=0, pad_mode='same',
                     weight_init=weight)


def _fc(in_channel, out_channel):
    weight_shape = (out_channel, in_channel)
    weight = Tensor(kaiming_uniform(weight_shape, a=math.sqrt(5)))
    return nn.Dense(in_channel, out_channel, has_bias=True, weight_init=weight, bias_init=0)


import mindspore.ops as ops


class ResidualBlock(nn.Cell):
    expansion = 4

    def __init__(self, in_channel, out_channel, stride=1, use_se=False, se_block=False):
        super(ResidualBlock, self).__init__()

        self.stride = stride
        self.use_se = use_se
        self.se_block = se_block
        chanel = out_channel // self.expansion
        self.conv1 = _conv1x1(in_channel, chanel, stride=1)
        self.bn1 = _bn(chanel)

        self.conv2 = _conv3x3(chanel, chanel, stride=stride)
        self.bn2 = _bn(chanel)

        self.conv3 = _conv1x1(chanel, out_channel, stride=1)
        self.bn3 = _bn(out_channel)
        self.relu = nn.ReLU()

        self.down_sample = False

        if stride != 1 or in_channel != out_channel:
            self.down_sample = True
        self.down_sample_layer = None

        if self.down_sample:
            self.down_sample_layer = nn.SequentialCell([_conv1x1(in_channel, out_channel, stride),
                                                        _bn(out_channel)])

    def construct(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.se_block:
            out_se = out
            out = self.se_global_pool(out, (2, 3))
            out = self.se_dense_0(out)
            out = self.relu(out)
            out = self.se_dense_1(out)
            out = self.se_sigmoid(out)
            out = ops.reshape(out, ops.shape(out) + (1, 1))
            out = self.se_mul(out, out_se)
        if self.down_sample:
            identity = self.down_sample_layer(identity)

        out = out + identity
        out = self.relu(out)

        return out


def _make_layer(self, block, layer_num, in_channel, out_channel, stride):
    layers = []
    resnet_block = block(in_channel, out_channel, stride=stride)
    layers.append(resnet_block)
    for _ in range(1, layer_num):
        resnet_block = block(out_channel, out_channel, stride=1)
        layers.append(resnet_block)
    return nn.SequentialCell(layers)


class ResNet(nn.Cell):
    def __init__(self, block, lay_nums, in_channels, out_channels, strides, num_classes, res_base=False):
        if not len(lay_nums) == len(in_channels) == len(out_channels) == 4:
            raise ValueError("the length of layer_num,in_channels,out_channels list must be 4")
        self.res_base = res_base

        self.conv1 = _con7x7(3, 64, stride=2, res_base=self.res_base)
        self.bn1 = _bn(64, self.res_base)
        self.relu = ops.relu()

        if self.res_base:
            self.pad = nn.Pad(paddings=((0, 0), (0, 0), (1, 1), (1, 1)))
            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='valid')
        else:
            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='same')

        self.layer1 = _make_layer(block, lay_nums[0], in_channel=in_channels[0], out_channel=out_channels[0],
                                  stride=strides[0])
        self.layer2 = _make_layer(block, lay_nums[1], in_channel=in_channels[1], out_channel=out_channels[1],
                                  stride=strides[1])
        self.layer3 = _make_layer(block, lay_nums[2], in_channel=in_channels[2], out_channel=out_channels[2],
                                  stride=strides[2])
        self.layer4 = _make_layer(block, lay_nums[3], in_channel=in_channels[3], out_channel=out_channels[3],
                                  stride=strides[3])

        self.mean = ops.ReduceMean(keep_dims=True)
        set.flatten = nn.Flatten()
        self.end_point = _fc(out_channels[3], num_classes)

    def construct(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        if self.res_base:
            x = self.pad(x)
        c1 = self.maxpool(x)

        c2 = self.layer1(c1)
        c3 = self.layer2(c2)
        c4 = self.layer3(c3)
        c5 = self.layer4(c4)

        out = self.mean(c5, (2, 3))
        out = self.flatten(out)
        out = self.end_point(out)

        return out


def resnet50(class_num=5):
    return ResNet(ResidualBlock, lay_nums=[3, 4, 6, 3], in_channels=[64, 256, 512, 1024],
                  out_channels=[256, 512, 1024, 2048],
                  strides=[1, 2, 2, 2], num_classes=class_num)


import glob
import mindspore.log as logger
import datetime

if cfg['pre_trained']:
    ckpt_parm_dict = ms.load_checkpoint(cfg['ckpt_path'])

data = create_dataset(dataset_path=os.path.join(cfg['data_path'], 'flowers', 'train'), batch_size=cfg['batch_size'],
                      do_train=True)

ds_val = create_dataset(dataset_path=os.path.join(cfg['data_path'], 'flowers', 'val'))

net = resnet50(class_num=cfg['num_classes'])
if cfg['pre_trained']:
    ms.load_param_into_net(net, ckpt_parm_dict)


def _generate_poly_lr(lr_init, lr_end, lr_max, total_steps, warmup_steps):
    lr_each_step = []
    if warmup_steps != 0:
        inc_each_step = (float(lr_max) - float(lr_init)) / float(warmup_steps)
    else:
        inc_each_step = 0
    for i in range(total_steps):
        if i < warmup_steps:
            lr = float(lr_init) + inc_each_step * float(i)
        else:
            base = (1.0 - (float(i) - float(warmup_steps)) / (float(total_steps) - float(warmup_steps)))
            lr = float(lr_max) * base * base
            if lr < 0.0:
                lr = 0.0
        lr_each_step.append(lr)
        return lr_each_step


import numpy as np


def get_lr(lr_init, lr_end, lr_max, warmupt_epochs, total_epochs, steps_per_epoch, lr_decay_mode):
    lr_each_step = []
    total_steps = steps_per_epoch * total_epochs
    warmupt_steps = steps_per_epoch * warmupt_epochs
    lr_each_step = _generate_poly_lr(lr_init, lr_end, lr_max, total_steps, warmupt_steps)
    lr_each_step = np.array(lr_each_step).astype(np.float32)
    return lr_each_step


def init_lr(step_size):
    if cfg['netname'] in ('resnet18', 'resnet34', 'resnet50', 'resnet152', 'se-resnet50',):
        lr = get_lr(lr_init=cfg['lr_init'], lr_end=cfg['lr_end'], lr_max=cfg['lr_max'],
                    warmupt_epochs=cfg['warmup_epochs']
                    , total_epochs=cfg['epoch_size'], steps_per_epoch=step_size, lr_decay_mode=cfg['lr_decay_mode'])
    return lr


step_size = 2
lr = ms.Tensor(init_lr(step_size=step_size))


def init_group_params(net):
    decayed_params = []
    no_decayed_params = []
    for parm in net.trainable_params():
        if 'beta' not in parm.name and 'gamma' not in parm.name and 'bias' not in parm.name:
            decayed_params.append(parm)
        else:
            no_decayed_params.append(parm)

    group_parms = [{'params': decayed_params, 'weight_decay': cfg['weight_decay']},
                   {'params': no_decayed_params}, {'order_params': net.trainable_params()}]
    return group_parms


group_params = init_group_params(net)
opt = nn.Momentum(group_params, lr, cfg['momentum'], loss_scale=cfg['loss_scale'])


def init_loss_scale():
    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    return loss


loss = init_loss_scale()
loss_scale = ms.FixedLossScaleManager(cfg['loss_scale'], drop_overflow_update=False)

metrics = {'acc'}
# if (cfg['netname'] in ('resnet18', 'resnet34', 'resnet50', 'resnet152', 'se-resnet50')) or contex.get